from read_data import parse_radar_data
import numpy as np
from scipy.interpolate import griddata
from typing import Dict, Tuple, List

def create_lat_lon_grid(parsed_data: Dict) -> Tuple[np.ndarray, np.ndarray]:
    """创建原始数据的经纬度网格"""
    # 获取四个角的经纬度
    (lon1, lat1), (lon2, lat2), (lon3, lat3), (lon4, lat4) = parsed_data['corners']

    # 创建经纬度网格（线性插值）
    lon_points = np.linspace(lon1, lon2, parsed_data['H'])
    lat_points = np.linspace(lat1, lat4, parsed_data['L'])

    return np.meshgrid(lon_points, lat_points)


def extract_region(parsed_data: Dict, min_lat: float, max_lat: float,
                   min_lon: float, max_lon: float) -> Dict:
    """提取指定经纬度范围内的数据"""
    # 创建经纬度网格
    lon_grid, lat_grid = create_lat_lon_grid(parsed_data)

    # 找到在目标范围内的索引
    mask = (lat_grid >= min_lat) & (lat_grid <= max_lat) & \
           (lon_grid >= min_lon) & (lon_grid <= max_lon)

    # 获取有效索引
    valid_rows, valid_cols = np.where(mask)

    if len(valid_rows) == 0:
        raise ValueError("指定范围内没有数据，请检查经纬度范围是否正确")

    # 提取边界索引
    min_row, max_row = np.min(valid_rows), np.max(valid_rows)
    min_col, max_col = np.min(valid_cols), np.max(valid_cols)

    # 提取子区域数据
    extracted = {
        'H': max_col - min_col + 1,
        'L': max_row - min_row + 1,
        'N': parsed_data['N'],
        'grid_size_km': parsed_data['grid_size_km'],
        'layer_heights': parsed_data['layer_heights'],
        'proj_mode': parsed_data['proj_mode'],
        'corners': [
            (lon_grid[min_row, min_col], lat_grid[min_row, min_col]),
            (lon_grid[min_row, max_col], lat_grid[min_row, max_col]),
            (lon_grid[max_row, max_col], lat_grid[max_row, max_col]),
            (lon_grid[max_row, min_col], lat_grid[max_row, min_col])
        ]
    }

    # 提取每个层的子区域数据
    extracted['data'] = []
    extracted['lon_grid'] = lon_grid[min_row:max_row + 1, min_col:max_col + 1]
    extracted['lat_grid'] = lat_grid[min_row:max_row + 1, min_col:max_col + 1]

    for layer in range(parsed_data['N']):
        extracted_layer = [
            row[min_col:max_col + 1]
            for row in parsed_data['data'][layer][min_row:max_row + 1]
        ]
        extracted['data'].append(extracted_layer)

    return extracted


def interpolate_to_1km(region_data: Dict, target_resolution_km: float = 1.0) -> Dict:
    """将数据插值到目标分辨率（默认1km×1km）"""
    # 计算目标网格的大小
    min_lon, max_lon = np.min(region_data['lon_grid']), np.max(region_data['lon_grid'])
    min_lat, max_lat = np.min(region_data['lat_grid']), np.max(region_data['lat_grid'])

    # 计算目标网格的格点数量
    # 基于原始网格的大致距离估算经纬度每度对应的公里数
    # 1度纬度约为111km，1度经度在38°N约为85km
    lon_per_km = 1 / 85.0  # longitude per km at ~38°N
    lat_per_km = 1 / 111.0  # latitude per km

    # 计算目标网格尺寸
    lon_size = max_lon - min_lon
    lat_size = max_lat - min_lat

    # 计算目标格点数量
    target_H = int(round(lon_size / (target_resolution_km * lon_per_km)))
    target_L = int(round(lat_size / (target_resolution_km * lat_per_km)))

    # 确保至少有一个格点
    target_H = max(1, target_H)
    target_L = max(1, target_L)

    # 创建目标网格
    target_lon = np.linspace(min_lon, max_lon, target_H)
    target_lat = np.linspace(min_lat, max_lat, target_L)
    target_lon_grid, target_lat_grid = np.meshgrid(target_lon, target_lat)

    # 准备插值数据
    points = np.column_stack([
        region_data['lon_grid'].flatten(),
        region_data['lat_grid'].flatten()
    ])

    # 插值每个层的数据
    interpolated_data = []
    for layer in range(region_data['N']):
        values = np.array(region_data['data'][layer]).flatten()
        # 使用线性插值
        interpolated_layer = griddata(
            points, values, (target_lon_grid, target_lat_grid),
            method='linear', fill_value=0.0  # 填充值可根据实际情况调整
        )
        interpolated_data.append(interpolated_layer)

    # 构建结果字典
    result = {
        'H': target_H,
        'L': target_L,
        'N': region_data['N'],
        'grid_size_km': target_resolution_km,
        'layer_heights': region_data['layer_heights'],
        'proj_mode': region_data['proj_mode'],
        'data': interpolated_data,
        'lon_grid': target_lon_grid,
        'lat_grid': target_lat_grid,
        'corners': [
            (min_lon, max_lat),
            (max_lon, max_lat),
            (max_lon, min_lat),
            (min_lon, min_lat)
        ]
    }

    return result